package core

import (
	"golang.org/x/net/bpf"
	"golang.org/x/net/ipv4"
)

func SetRawConnBPF(rc *ipv4.RawConn, filter []bpf.Instruction) error {
	assembled, err := bpf.Assemble(filter)
	if err != nil {
		return err
	}
	return rc.SetBPF(assembled)
}

func CreateBPFDropAll() []bpf.Instruction {
	return []bpf.Instruction{
		bpf.RetConstant{Val: 0x0},
	}
}

func SetBPFDropAll(rc *ipv4.RawConn) error {
	filter := CreateBPFDropAll()
	assembled, err := bpf.Assemble(filter)
	if err != nil {
		return err
	}
	return rc.SetBPF(assembled)
}

func CreateBPFICMP(id uint16, receiverID uint16) []bpf.Instruction {
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                 // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(48), SkipTrue: 12}, // ip header + icmp + payload

		bpf.LoadMemShift{Off: 0},                                      // skip ip header
		bpf.LoadIndirect{Off: 0, Size: 1},                             // icmp type,
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(0), SkipFalse: 9}, // echo reply

		bpf.LoadIndirect{Off: 4, Size: 2},                              // echo reply identifier
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(id), SkipFalse: 7}, // echo reply

		bpf.LoadIndirect{Off: 8, Size: 4},                              // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x6e70726f, SkipFalse: 5}, // npro

		bpf.LoadIndirect{Off: 12, Size: 4},                             // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x62652121, SkipFalse: 3}, // be!!

		bpf.LoadIndirect{Off: 16, Size: 2},                                     // receiverID
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(receiverID), SkipFalse: 1}, // receiverID

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFUDPICMP(srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                 // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(56), SkipTrue: 20}, // ip header + icmp + ip header + udp

		bpf.LoadMemShift{Off: 0},                                       // skip ip header
		bpf.LoadIndirect{Off: 0, Size: 1},                              // icmp type,
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(3), SkipFalse: 17}, // destination unreachable

		bpf.LoadIndirect{Off: 17, Size: 1},                              // ip protocol
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(17), SkipFalse: 15}, // udp

		bpf.TXA{},
		bpf.ALUOpConstant{Op: bpf.ALUOpAdd, Val: 8},
		bpf.TAX{},
		bpf.LoadIndirect{Off: 0, Size: 1}, // udp length
		bpf.ALUOpConstant{Op: bpf.ALUOpAnd, Val: 0x0F}, // Mask to get IHL (low 4 bits)
		bpf.ALUOpConstant{Op: bpf.ALUOpShiftLeft, Val: 2},
		bpf.ALUOpX{Op: bpf.ALUOpAdd},
		bpf.TAX{},

		bpf.LoadIndirect{Off: 0, Size: 2},                                               // udp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 5}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 4},    // max

		bpf.LoadIndirect{Off: 2, Size: 2},                                               // udp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 2}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 1},    // max

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFUDPICMP2(srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                 // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(56), SkipTrue: 12}, // ip header + icmp + ip header + udp

		bpf.LoadMemShift{Off: 0},                                      // skip ip header
		bpf.LoadIndirect{Off: 0, Size: 1},                             // icmp type,
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(3), SkipFalse: 9}, // destination unreachable

		bpf.LoadIndirect{Off: 17, Size: 1},                             // ip protocol
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(17), SkipFalse: 7}, // udp

		// bpf.LoadMemShift{Off: 8}, // skip ip header

		bpf.LoadIndirect{Off: 28, Size: 2},                                              // udp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 5}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 4},    // max

		bpf.LoadIndirect{Off: 30, Size: 2},                                              // udp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 2}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 1},    // max

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFUDP(pid uint16, receiverID uint16, srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	return []bpf.Instruction{

		bpf.LoadMemShift{Off: 0}, // skip ip header

		bpf.LoadIndirect{Off: 0, Size: 2},                                                // udp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 15}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 14},    // max

		bpf.LoadIndirect{Off: 2, Size: 2},                                                // udp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 12}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 11},    // max

		bpf.LoadIndirect{Off: 4, Size: 2},                                // udp length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(20), SkipTrue: 9}, // 20 bytes

		bpf.LoadIndirect{Off: 8, Size: 4},                              // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x6e70726f, SkipFalse: 7}, // npro

		bpf.LoadIndirect{Off: 12, Size: 4},                             // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x62652121, SkipFalse: 5}, // be!!

		bpf.LoadIndirect{Off: 16, Size: 2},                                     // receiverID
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(receiverID), SkipFalse: 3}, // receiverID

		bpf.LoadIndirect{Off: 18, Size: 2},                              // pid
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(pid), SkipFalse: 1}, // 7

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFUDPServer(srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(28), SkipTrue: 8}, // ip header + udp + payload

		bpf.LoadMemShift{Off: 0}, // skip ip header

		bpf.LoadIndirect{Off: 0, Size: 2},                                               // udp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 5}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 4},    // max

		bpf.LoadIndirect{Off: 2, Size: 2},                                               // udp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 2}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 1},    // max

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFTCPSYN(srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	tcpFlagAck := 0x10
	tcpFlagRst := 0x04
	tcpFlagSynAck := 0x12
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                 // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(40), SkipTrue: 12}, // ip header + tcp header

		bpf.LoadMemShift{Off: 0}, // skip ip header

		bpf.LoadIndirect{Off: 0, Size: 2},                                               // tcp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 9}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 8},    // max

		bpf.LoadIndirect{Off: 2, Size: 2},                                               // tcp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 6}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 5},    // max

		bpf.LoadIndirect{Off: 13, Size: 1},                                          // tcp flag
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(tcpFlagAck), SkipTrue: 2},       // ACK包
		bpf.JumpIf{Cond: bpf.JumpBitsSet, Val: uint32(tcpFlagRst), SkipTrue: 1},     // RST包,
		bpf.JumpIf{Cond: bpf.JumpBitsSet, Val: uint32(tcpFlagSynAck), SkipFalse: 1}, // SYN+ACK包

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFTCP(pid uint16, receiverID uint16, srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                 // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(62), SkipTrue: 16}, // ip header + tcp + payload

		bpf.LoadMemShift{Off: 0}, // skip ip header

		bpf.LoadIndirect{Off: 0, Size: 2},                                                // tcp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 13}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 12},    // max

		bpf.LoadIndirect{Off: 2, Size: 2},                                                // tcp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 10}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 9},     // max

		bpf.LoadIndirect{Off: 20, Size: 4},                             // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x6e70726f, SkipFalse: 7}, // npro

		bpf.LoadIndirect{Off: 24, Size: 4},                             // magic number
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: 0x62652121, SkipFalse: 5}, // be!!

		bpf.LoadIndirect{Off: 28, Size: 2},                                     // receiverID
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(receiverID), SkipFalse: 3}, // receiverID

		bpf.LoadIndirect{Off: 30, Size: 2},                              // pid
		bpf.JumpIf{Cond: bpf.JumpEqual, Val: uint32(pid), SkipFalse: 1}, // 7

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}

func CreateBPFTCPServer(srcPortMin uint16, srcPortMax uint16, dstPortMin uint16, dstPortMax uint16) []bpf.Instruction {
	tcpFlagRst := 0x04
	tcpFlagSyn := 0x02
	return []bpf.Instruction{
		bpf.LoadIndirect{Off: 2, Size: 2},                                 // ip total length
		bpf.JumpIf{Cond: bpf.JumpLessThan, Val: uint32(40), SkipTrue: 11}, // ip header + tcp + payload

		bpf.LoadMemShift{Off: 0}, // skip ip header

		bpf.LoadIndirect{Off: 0, Size: 2},                                               // tcp src port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(srcPortMin), SkipFalse: 8}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(srcPortMax), SkipFalse: 7},    // max

		bpf.LoadIndirect{Off: 2, Size: 2},                                               // tcp dst port
		bpf.JumpIf{Cond: bpf.JumpGreaterOrEqual, Val: uint32(dstPortMin), SkipFalse: 5}, // min
		bpf.JumpIf{Cond: bpf.JumpLessOrEqual, Val: uint32(dstPortMax), SkipFalse: 4},    // max

		bpf.LoadIndirect{Off: 13, Size: 1},                                      // tcp flag
		bpf.JumpIf{Cond: bpf.JumpBitsSet, Val: uint32(tcpFlagRst), SkipTrue: 2}, // RST包,
		bpf.JumpIf{Cond: bpf.JumpBitsSet, Val: uint32(tcpFlagSyn), SkipTrue: 1}, // SYN包

		bpf.RetConstant{Val: 65535}, // capture
		bpf.RetConstant{Val: 0x0},   // drop
	}
}
